from typing import Union

import yaml
import fire
import tqdm
import os
import sys
from pathlib import Path

sys.path.append(str(Path(__file__).absolute().parent))
import numpy as np
import torch
import torchaudio
import torch.multiprocessing as mp

from model import ConvTasNet
from utils.load_scp import read_2column_text

def normalize(raw_wav, est_wav):
    norm = torch.norm(raw_wav, p=float('inf'))
    est_wav = est_wav * norm / torch.max(torch.abs(est_wav))
    return est_wav

def run(rank, world_size, config, ckpt, mix_scp, ref_scp, output_dir, device):
    device = device if isinstance(device, str) else device[rank % len(device)]
    with open(config, "r") as f:
        config = yaml.safe_load(f)
    with open(config['spk_id'], "r") as f:
        spk_len = len(f.readlines())
    model = ConvTasNet(**config['model_conf'], num_spks=spk_len)
    ckpt = torch.load(ckpt, map_location='cpu')
    model.load_state_dict(ckpt['model_state_dict'])
    model.eval()
    model.to(device)
    mix_scp_dict = read_2column_text(mix_scp)
    ref_scp_dict = read_2column_text(ref_scp)
    keys = list(mix_scp_dict.keys())
    keys = keys[rank::world_size]
    with torch.no_grad():
        for _k in tqdm.tqdm(keys, desc=f"[rank {rank}]"):
            mix_path = mix_scp_dict[_k]
            ref_path = ref_scp_dict[_k]
            mix, sr = torchaudio.load(mix_path) 
            ref, sr = torchaudio.load(ref_path) # [1,T]
            mix, ref = mix.to(device), ref.to(device)
            ref_len = torch.full((ref.size(0),), ref.size(1)).to(ref.device) # [B]
            audio_hat, _, _, _ = model(mix, ref, ref_len) #[1,T]
            audio_hat = normalize(mix, audio_hat)

            
            save_path = str(Path(output_dir) / f"{Path(mix_path).stem}.wav")
            torchaudio.save(save_path, audio_hat.cpu(), sample_rate=sr)


def main(device: Union[list, str], num_proc:int, config, ckpt, mix_scp, ref_scp, output_dir):
    """
    
    """
    os.makedirs(output_dir, exist_ok=True)
    if num_proc > 1:
        print(f"running {num_proc} processes")
        mp.spawn(run, nprocs = num_proc, args=(num_proc, config, ckpt, mix_scp, ref_scp, output_dir, device))
    else:
        run(0, 1, num_proc, config, ckpt, mix_scp, ref_scp, output_dir, device)
    print("Done...")


if __name__ == "__main__":
    fire.Fire(main)
    pass